#!/usr/bin/env python3
import os
import time
import subprocess
import signal
import sys

def create_veth_pairs():
    """创建虚拟以太网接口对"""
    print("📡 Creating virtual Ethernet interfaces...")

    # 创建所有需要的 veth 对
    # h1-s1: veth1-veth2
    # s1-s2: veth3-veth4
    # s1-s3: veth5-veth6
    # s1-s4: veth7-veth8
    # s1-s5: veth9-veth10
    # s2-s6: veth11-veth12
    # s3-s6: veth13-veth14
    # s4-s6: veth15-veth16
    # s5-s6: veth17-veth18
    # s6-h2: veth19-veth20

    for i in range(1, 21):
        if i % 2 == 1:  # 只创建奇数编号的接口，偶数会自动创建
            cmd = f'sudo ip link add name veth{i} type veth peer name veth{i+1}'
            subprocess.run(cmd, shell=True)

    # 启动所有接口
    for i in range(1, 21):
        subprocess.run(f'sudo ip link set dev veth{i} up', shell=True)

def start_p4_switches():
    """启动 P4 交换机"""
    print("🚀 Starting P4 switches...")

    switches = [
        {
            'name': 's1', 'port': 9091,
            'interfaces': ['1@veth1', '2@veth3', '3@veth5', '4@veth7', '5@veth9']
        },
        {
            'name': 's2', 'port': 9092,
            'interfaces': ['1@veth4', '2@veth11']
        },
        {
            'name': 's3', 'port': 9093,
            'interfaces': ['1@veth6', '2@veth13']
        },
        {
            'name': 's4', 'port': 9094,
            'interfaces': ['1@veth8', '2@veth15']
        },
        {
            'name': 's5', 'port': 9095,
            'interfaces': ['1@veth10', '2@veth17']
        },
        {
            'name': 's6', 'port': 9096,
            'interfaces': ['1@veth19', '2@veth12', '3@veth14', '4@veth16', '5@veth18']
        }
    ]

    processes = []

    for switch in switches:
        interface_args = ' '.join([f'-i {iface}' for iface in switch['interfaces']])
        cmd = f'simple_switch_grpc --thrift-port {switch["port"]} --device-id {switches.index(switch)} {interface_args} build/load_balance.json > logs/{switch["name"]}.log 2>&1 &'

        print(f"Starting {switch['name']} on thrift port {switch['port']}...")
        result = subprocess.run(cmd, shell=True)
        processes.append(switch['name'])

    return processes

def setup_hosts():
    """配置主机网络命名空间"""
    print("🖥️ Setting up host network namespaces...")

    # 创建 h1 命名空间
    subprocess.run('sudo ip netns add h1', shell=True)
    subprocess.run('sudo ip link set veth2 netns h1', shell=True)
    subprocess.run('sudo ip netns exec h1 ip link set veth2 name eth0', shell=True)
    subprocess.run('sudo ip netns exec h1 ip addr add 10.0.1.1/24 dev eth0', shell=True)
    subprocess.run('sudo ip netns exec h1 ip link set eth0 up', shell=True)
    subprocess.run('sudo ip netns exec h1 ip route add default via 10.0.1.10', shell=True)
    subprocess.run('sudo ip netns exec h1 arp -s 10.0.1.10 08:00:00:00:01:00', shell=True)

    # 创建 h2 命名空间
    subprocess.run('sudo ip netns add h2', shell=True)
    subprocess.run('sudo ip link set veth20 netns h2', shell=True)
    subprocess.run('sudo ip netns exec h2 ip link set veth20 name eth0', shell=True)
    subprocess.run('sudo ip netns exec h2 ip addr add 10.0.2.2/24 dev eth0', shell=True)
    subprocess.run('sudo ip netns exec h2 ip link set eth0 up', shell=True)
    subprocess.run('sudo ip netns exec h2 ip route add default via 10.0.2.20', shell=True)
    subprocess.run('sudo ip netns exec h2 arp -s 10.0.2.20 08:00:00:00:02:00', shell=True)

def run_host_command(host, command):
    """在主机命名空间中运行命令"""
    cmd = f'sudo ip netns exec {host} {command}'
    return subprocess.run(cmd, shell=True, capture_output=True, text=True)

def cleanup():
    """清理函数"""
    print("🧹 Cleaning up...")
    subprocess.run('sudo pkill -f simple_switch_grpc', shell=True)
    subprocess.run('sudo ip netns del h1 2>/dev/null', shell=True)
    subprocess.run('sudo ip netns del h2 2>/dev/null', shell=True)

    # 删除虚拟接口
    for i in range(1, 21):
        subprocess.run(f'sudo ip link delete veth{i} 2>/dev/null', shell=True)

def signal_handler(sig, frame):
    """处理 Ctrl+C 信号"""
    print('\n🛑 Received Ctrl+C, shutting down...')
    cleanup()
    sys.exit(0)

def main():
    # 设置信号处理
    signal.signal(signal.SIGINT, signal_handler)

    # 创建目录
    os.makedirs('logs', exist_ok=True)
    os.makedirs('pcaps', exist_ok=True)

    print("🎯 Starting Pure P4 ECMP Load Balancing Topology")
    print("=" * 50)

    # 清理之前的实例
    cleanup()
    time.sleep(2)

    try:
        # 步骤 1: 创建虚拟接口
        create_veth_pairs()
        time.sleep(1)

        # 步骤 2: 配置主机
        setup_hosts()
        time.sleep(1)

        # 步骤 3: 启动 P4 交换机
        processes = start_p4_switches()

        print("⏳ Waiting for switches to initialize (15 seconds)...")
        time.sleep(15)

        print("\n" + "=" * 60)
        print("✅ ECMP Load Balancing Topology is READY!")
        print("=" * 60)
        print("Topology: h1 ↔ s1 ↔ [s2,s3,s4,s5] ↔ s6 ↔ h2")
        print("Forward ECMP: 10.0.0.1 → 4 paths (h1→h2)")
        print("Reverse ECMP: 10.0.0.2 → 4 paths (h2→h1)")

        print("\n📋 Next steps:")
        print("1. Open NEW terminal and run: make controller")
        print("2. Test connectivity:")
        print("   - In h1: sudo ip netns exec h1 ping 10.0.2.2")
        print("   - Test ECMP: sudo ip netns exec h1 ./send.py 10.0.0.1 'Test Packet'")
        print("3. Monitor traffic: sudo tcpdump -i veth3 -c 5 -nn -q &")

        print("\n🔄 Running interactive shell...")
        print("Type 'exit' to shutdown the topology")

        # 保持运行
        while True:
            try:
                user_input = input("p4-topology> ").strip().lower()
                if user_input in ['exit', 'quit', 'q']:
                    break
                elif user_input == 'status':
                    print("✅ Topology is running")
                elif user_input == 'test':
                    print("Testing connectivity...")
                    result = run_host_command('h1', 'ping -c 3 10.0.2.2')
                    print(result.stdout)
                elif user_input == 'help':
                    print("Commands: status, test, exit")
                else:
                    print("Unknown command. Type 'help' for available commands.")
            except EOFError:
                break

    except Exception as e:
        print(f"❌ Error: {e}")
    finally:
        cleanup()
        print("✅ Cleanup complete!")

if __name__ == '__main__':
    main()
